#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""本文件为示例主程序，完整说明请见同目录的 README.md。"""


import os
import re
import math
from typing import List, Dict, Any, Tuple, Optional
from dataclasses import dataclass, field
from collections import deque
from heapq import nlargest
from time import time

# 仅尝试加载 .env（如果没有安装 python-dotenv，不会报错；此时需自行导出环境变量）
try:
    from dotenv import load_dotenv  # type: ignore
    load_dotenv(".env")
except Exception:
    pass

from zhipuai import ZhipuAI

# =========================
# 工具函数（通用）
# =========================

def now_ts() -> float:
    """返回当前时间戳（秒）。用于给消息/向量条目打时间标记。"""
    return time()

def tokenize(text: str) -> List[str]:
    """
    极简分词器：
    - 仅用于估算“词数≈token数”，帮助做 STM 的“词数预算”裁剪；
    - 规则：小写化后，提取字母/数字/中文统一表意文字，避免符号/空白的干扰。
    """
    return re.findall(r"[a-zA-Z0-9\u4e00-\u9fff]+", text.lower())

def cosine(a: List[float], b: List[float]) -> float:
    """
    余弦相似度：衡量两个向量“方向”相似度，范围 [-1,1]（本例均为非负）。
    - 在向量检索中广泛使用；本例中的向量由智谱 Embeddings 产生。
    """
    if not a or not b or len(a) != len(b):
        return 0.0
    dot = sum(x*y for x, y in zip(a, b))
    na = math.sqrt(sum(x*x for x in a)) or 1.0
    nb = math.sqrt(sum(y*y for y in b)) or 1.0
    return dot / (na * nb)

# =========================
# STM：短期记忆（只保留最近上下文 + 简要摘要）
# =========================

@dataclass
class Message:
    """对话消息的轻量结构。role 取值：'user'/'assistant'/'system'。"""
    role: str
    content: str
    ts: float = field(default_factory=now_ts)

class ShortTermMemory:
    """
    目标：始终保持“够用”的最近上下文，且总体大小受控。
    做法：
      - 限制最大回合数（max_turns）；
      - 用“词数预算”（max_words）估算上下文体积，若超限→生成极简“摘要”，并裁掉一半旧消息。
    说明：
      - 这里的“摘要”是启发式实现（取“首句”拼接），为了示例尽量不引入 LLM。
      - 如果你追求摘要质量，可把 _naive_summary 改为“调用 LLM 生成摘要”。
    """
    def __init__(self, max_turns: int = 12, max_words: int = 800):
        self.turns: deque[Message] = deque()
        self.max_turns = max_turns
        self.max_words = max_words
        self._summary: Optional[str] = None

    def append(self, role: str, content: str):
        """追加一条消息到短期记忆，并在必要时触发裁剪/摘要。"""
        self.turns.append(Message(role=role, content=content))
        # 先按“回合数”粗裁
        while len(self.turns) > self.max_turns:
            self.turns.popleft()
        # 再按“词数预算”细裁
        self._trim_if_needed()

    def _trim_if_needed(self):
        """若当前上下文超出词数预算：先摘要“将被丢弃的旧消息”，再删除它们。"""
        text = self.render(include_summary=False)
        if len(tokenize(text)) > self.max_words:
            # 先锁定要丢弃的旧半区，并对其做摘要，避免信息完全丢失
            half = len(self.turns) // 2
            to_discard = list(self.turns)[:half]
            self._summary = self._naive_summary(to_discard)
            # 丢弃一半最早消息，保留“近期上下文”
            for _ in range(half):
                self.turns.popleft()

    def _naive_summary(self, msgs: Optional[List[Message]] = None) -> str:
        """
        极简摘要策略：
        - 若传入 msgs，则对 msgs 全量抽取“第一句”拼摘要；
        - 否则取 self.turns 的前 N 条（≤6），抽取“第一句”拼摘要。
        - 这样做能在完全不依赖 LLM 的情况下，给上下文一个“压缩版线索”。
        """
        corpus = msgs if msgs is not None else list(self.turns)[: min(6, len(self.turns))]
        sents = []
        for m in corpus:
            s0 = re.split(r"[。.!?！？]", m.content.strip())
            if s0 and s0[0]:
                sents.append(f"{m.role}: {s0[0]}")
        return " | ".join(sents)

    def render(self, include_summary: bool = True) -> str:
        """
        把当前 STM 转成字符串：
        - 若 include_summary=True 且存在摘要，则先拼摘要；
        - 然后按顺序输出各条消息。
        """
        chunks = []
        if include_summary and self._summary:
            chunks.append(f"[STM-SUMMARY] {self._summary}")
        for m in self.turns:
            chunks.append(f"{m.role}: {m.content}")
        return "\n".join(chunks)

    @property
    def summary(self) -> str:
        """返回当前的极简摘要文本（可能为空字符串）。"""
        return self._summary or ""

    def clear(self):
        """清空短期记忆及摘要，用于“强制重置上下文”的场景。"""
        self.turns.clear()
        self._summary = None

# =========================
# LTM：向量库（智谱 Embeddings + 纯内存检索）
# =========================

class ZhipuEmbedder:
    """
    智谱 Embeddings 包装：
    - embed(texts) 接收字符串列表，返回对应的向量列表（二维数组）。
    - 向量维度由模型决定（embedding-2 / embedding-3 等）。
    """
    def __init__(self, client: ZhipuAI, model: str = "embedding-2"):
        self.client = client
        self.model = model

    def embed(self, texts: List[str]) -> List[List[float]]:
        resp = self.client.embeddings.create(model=self.model, input=texts)
        return [item.embedding for item in resp.data]

@dataclass
class VSItem:
    """
    向量库条目：
    - id：内部递增ID；
    - text：原文片段（本例存“用户说的话”或知识点）；
    - vector：对应的向量表示；
    - meta：附加元信息（如 salience）；
    - ts：时间戳（可用于“新鲜度”策略）。
    """
    id: str
    text: str
    vector: List[float]
    meta: Dict[str, Any] = field(default_factory=dict)
    ts: float = field(default_factory=now_ts)

class InMemoryVectorStore:
    """
    纯内存向量库（最小实现）：
      - add(text, vector)：添加条目
      - search(query_vec, k, min_score)：基于余弦相似度做TopK检索
    说明：
      - 生产可替换为 FAISS/pgvector/ES 等，接口保持一致即可。
    """
    def __init__(self):
        self.items: Dict[str, VSItem] = {}
        self._id = 0

    def add(self, text: str, vector: List[float], meta: Optional[Dict[str, Any]] = None) -> str:
        self._id += 1
        _id = f"mem_{self._id}"
        self.items[_id] = VSItem(id=_id, text=text, vector=vector, meta=meta or {})
        return _id

    def search(self, query_vec: List[float], k: int = 4, min_score: float = 0.1) -> List[Tuple[VSItem, float]]:
        """
        相似检索：
        - 输入：query_vec（查询向量），k（返回条数），min_score（过滤阈值）；
        - 输出：[(VSItem, score), ...]，按相似度降序；
        - 这里只做最简单的线性扫描，足够示例用。
        """
        scored = ((it, cosine(query_vec, it.vector)) for it in self.items.values())
        topk = [p for p in nlargest(k, scored, key=lambda p: p[1]) if p[1] >= min_score]
        return topk

# =========================
# 是否写入 LTM：一个最小的启发式（salience）
# =========================

# 若用户话中出现这些提示词，则更倾向“要记住”
HINT_WORDS = ["记住", "我的偏好", "我喜欢", "不要忘", "prefer", "i like", "my preference", "remember"]

def estimate_novelty(store: InMemoryVectorStore, embedder: ZhipuEmbedder, text: str) -> float:
    """
    新颖度估计：1 - 与“库内最相似条目”的相似度
    - 若库为空 → 新颖度=1.0（全新信息）；
    - 否则：embed(text) → search Top1 → 取 best_score → novelty = 1 - best_score；
    - 该数值用于避免重复存储“差不多”的内容。
    """
    if not store.items:
        return 1.0
    qv = embedder.embed([text])[0]
    hits = store.search(qv, k=1, min_score=0.0)
    best = hits[0][1] if hits else 0.0
    return max(0.0, 1.0 - best)

def important_hint(text: str) -> float:
    """是否出现“请记住/偏好”等关键词（出现则返回1.0，否则0.0）。"""
    t = text.lower()
    return 1.0 if any(h in t for h in HINT_WORDS) else 0.0

def compute_salience(store: InMemoryVectorStore, embedder: ZhipuEmbedder, user_msg: str) -> float:
    """
    综合打分（0..1）：salience = 0.7 * 新颖度 + 0.3 * 关键词提示
    - 分数高，越“值得写入长期记忆”。
    - 你可以按业务调整权重/规则（如加入用户显式“标记记住”的指令）。
    """
    lower_msg = user_msg.lower()
    question = bool(re.search(r"[?？]", user_msg))
    explicit_remember = any(k in lower_msg for k in ["记住", "不要忘", "remember"])
    hint = important_hint(user_msg)
    # 问句默认不写入 LTM；只有显式要求“记住/不要忘/remember”时才保留
    if question and not explicit_remember:
        return 0.0
    # 打招呼等超短消息且无提示词时，不写入 LTM
    if not hint and len(tokenize(user_msg)) < 4:
        return 0.0
    novelty = estimate_novelty(store, embedder, user_msg)
    return max(0.0, min(1.0, 0.7*novelty + 0.3*hint))

# =========================
# MemoryManager：把 STM 与 LTM 粘起来
# =========================

class MemoryManager:
    """
    对外提供三个能力：
      1) observe_user / observe_agent：把每一轮的 user/assistant 消息纳入 STM；
      2) retrieve_for：给定 query，返回“STM 文本 + LTM 命中TopK”，用于拼提示词；
      3) maybe_persist：依据 salience 决定是否把用户消息写入 LTM。
    """
    def __init__(self, embedder: ZhipuEmbedder, store: InMemoryVectorStore,
                 stm_max_turns: int = 12, stm_max_words: int = 800):
        self.embedder = embedder
        self.store = store
        self.stm = ShortTermMemory(stm_max_turns, stm_max_words)

    def observe_user(self, text: str):
        """把用户消息塞进 STM。"""
        self.stm.append("user", text)

    def observe_agent(self, text: str):
        """把助手回复塞进 STM。"""
        self.stm.append("assistant", text)

    def maybe_persist(self, user_msg: str, threshold: float = 0.55) -> Optional[str]:
        """
        若 salience ≥ 阈值，则把用户消息写入 LTM。
        - 返回：写入后的条目ID；若不满足阈值则返回 None。
        """
        sal = compute_salience(self.store, self.embedder, user_msg)
        if sal >= threshold:
            vec = self.embedder.embed([user_msg])[0]
            return self.store.add(user_msg, vec, meta={"salience": sal})
        return None

    def retrieve_for(self, query: str, k: int = 4) -> Dict[str, Any]:
        """
        根据当前问题检索上下文：
        - 把 STM 的摘要拼在 query 后（能提升“与当前对话语境一致”的召回）；
        - 用向量检索拿到 LTM TopK 命中；
        - 返回 { "stm_text": ..., "ltm_hits": [{"text":..., "score":...}, ...] }
        """
        q = query if not self.stm.summary else f"{query} || {self.stm.summary}"
        qv = self.embedder.embed([q])[0]
        hits = self.store.search(qv, k=k)
        return {
            "stm_text": self.stm.render(),
            "ltm_hits": [{"text": it.text, "score": float(sc)} for it, sc in hits]
        }

# =========================
# LLM：智谱 Chat Completions 封装
# =========================

class ZhipuChat:
    """
    对 zhipuai SDK 的轻封装：
    - ask(messages) 按 OpenAI 风格传 messages（system/user/assistant），返回字符串回答；
    - 通过 extra_body 传入 temperature / max_tokens 等生成参数。
    """
    def __init__(self, client: ZhipuAI, model: str = "glm-4", temperature: float = 0.2, max_tokens: int = 512):
        self.client = client
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens

    def ask(self, messages: List[Dict[str, Any]]) -> str:
        resp = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            extra_body={"temperature": self.temperature, "max_tokens": self.max_tokens}
        )
        return resp.choices[0].message.content

# =========================
# 提示词组装：把 STM/LTM 注入到 system prompt
# =========================

def build_system_prompt(ctx: Dict[str, Any]) -> str:
    """
    把“短期对话摘要 + 长期记忆TopK”拼进 system prompt：
    - 让 LLM “知道你说过什么 & 你曾经的偏好/事实”；
    - 仍要求 LLM 结合常识与当前问题谨慎作答。
    """
    stm = ctx.get("stm_text", "")
    ltm = ctx.get("ltm_hits", [])
    refs = "\n".join([f"- ({h['score']:.2f}) {h['text']}" for h in ltm]) or "（无命中）"
    return (
        "你是一名中文助理。请在回答时参考“记忆/资料”，但不要生搬硬套；"
        "如与用户问题或常识冲突，以用户问题为主。回答要简洁、可执行。\n\n"
        f"【短期对话摘要】\n{stm}\n\n"
        f"【长期记忆 Top{len(ltm)}】\n{refs}\n"
    )

# =========================
# 演示主程序：一组最小的对话回合
# =========================

def run_demo():
    # 1) 读取必要配置
    api_key = os.getenv("ZHIPUAI_API_KEY", "").strip()
    if not api_key:
        raise RuntimeError("未找到 ZHIPUAI_API_KEY，请在 .env 或系统环境中设置。")
    base_url = os.getenv("ZHIPUAI_BASE_URL")  # 可选
    llm_model = os.getenv("ZHIPU_LLM_MODEL", "glm-4.5")
    emb_model = os.getenv("ZHIPU_EMBED_MODEL", "embedding-2")

    # 2) 初始化 ZhipuAI 客户端与组件
    client = ZhipuAI(api_key=api_key, base_url=base_url) if base_url else ZhipuAI(api_key=api_key)
    chat = ZhipuChat(client, model=llm_model, temperature=0.2, max_tokens=512)
    embedder = ZhipuEmbedder(client, model=emb_model)
    store = InMemoryVectorStore()
    mm = MemoryManager(embedder, store, stm_max_turns=10, stm_max_words=200)

    # 3) 示例对话分两段，中间强制清空 STM（LTM 保留）
    dialogues = [
        [
            ("user", "你好！"),
            ("user", "我喜欢乌龙茶，不喜欢太甜的饮料，请记住。"),
            ("user", "推荐一杯饮料？"),
        ],
        [
            ("user", "昨天我说我喜欢什么茶？"),
            ("user", "给个不含糖的下午茶建议。"),
        ],
    ]

    for idx, dialogue in enumerate(dialogues):
        if idx > 0:
            mm.stm.clear()
            print("\n[STM] 已清空（分段重置上下文）")
        for role, content in dialogue:
            if role == "user":
                # ① 记录用户输入进 STM
                print(f"\n[user] {content}")
                mm.observe_user(content)
                print(f"[STM] 已记录 user 消息，共 {len(mm.stm.turns)} 条")

                # ② 基于“当前问题+STM摘要”从 LTM 检索相关记忆，构造 system prompt
                ctx = mm.retrieve_for(content, k=3)
                sys_prompt = build_system_prompt(ctx)
                messages = [
                    {"role": "system", "content": sys_prompt},
                    {"role": "user", "content": content},
                ]

                # ③ 询问 LLM 并打印回复，同时把回复也纳入 STM（便于后续问答）
                answer = chat.ask(messages)
                print(f"[assistant]\n{answer}")
                mm.observe_agent(answer)
                print(f"[STM] 已记录 assistant 消息，共 {len(mm.stm.turns)} 条")

                # ④ 尝试把“本轮用户话”写入 LTM（若 salience 达阈值）
                persisted = mm.maybe_persist(content, threshold=0.55)
                if persisted:
                    print(f"[LTM] 已写入：{persisted}")
            else:
                # 固定助手语句也纳入 STM，以保持完整上下文
                print(f"\n[assistant] {content}")
                mm.observe_agent(content)

if __name__ == "__main__":
    run_demo()
